# ==============================================================================
# Copyright (c) 2025 CompanyName. All rights reserved.
# Author:         22020873 陈泽欣
# Project:        Design of Deep Learning Fundamental Course
# Module:         config.py
# Date:           2025-05-24
# Description:    此模块负责深度学习模型的可视化过程，包括但不限于模型初始化、图像处理参数设置、数据记录配置等。
#                 它集成了PoseNet和YOLO模型，用于特定对象的姿态估计和识别。通过该脚本，用户可以加载预训练模型，
#                 并对输入图像进行处理，以实现高精度的对象检测与姿态估计。此外，还提供了一系列可调参数，便于用户
#                 根据实际应用场景优化算法性能。
# ==============================================================================

import os
import torch
from ultralytics import YOLO
from core.posenet_detect import PoseNet
from utils.utils import ensure_file
import cv2

class Config:
    # +-----------------------------------------------------------------+
    # 📦 1. 模型相关配置（Model Configuration）
    # 初始化模型并加载权重
    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

    # 加载 PoseNet 模型
    PoseNet_MODEL = PoseNet().to(DEVICE)
    PoseNet_MODEL.load_state_dict(torch.load('models\\dice_PoseNet.pth', map_location=DEVICE))
    print("✅ PoseNet_Model loaded successfully.")

    # 加载 YOLO 模型
    YOLO_MODEL = YOLO('models\\dice_yoloBestV3.pt')
    print("✅ YOLO_Model loaded successfully.")
    # +-----------------------------------------------------------------+


    # +-----------------------------------------------------------------+
    # 📁 2. 数据记录与保存配置（Data Saving）
    # 是否启用数据收集训练标志
    SAVE_DATA_TO_TRAIN_FLAG = False

    # 保存训练数据 CSV 
    SAVE_DATA_TO_TRAIN_CSV_FILE = "data\\poseDatas.csv"

    @classmethod
    def ensure_data_path(cls):
        ensure_file(cls.SAVE_DATA_TO_TRAIN_CSV_FILE)

    IMG_COUNT = len(os.listdir('data\\reproImages'))
    # +-----------------------------------------------------------------+


    # +-----------------------------------------------------------------+
    # 🔍 3. 置信度与误差阈值（Confidence and Error Thresholds）
    # YOLO检测置信度阈值，低于此值的检测结果被忽略
    ALLOW_YOLO_CONFIDENCE_THRESHOLD = 0.9

    # 不同方法下的最小允许重投影误差（越低越严格）
    NORMAL_MIN_ALLOW_REPROJECTION_ERROR = 5
    POSENET_MIN_ALLOW_REPROJECTION_ERROR = 15
    FUSION_MIN_ALLOW_REPROJECTION_ERROR = 5
    # +-----------------------------------------------------------------+


    # +-----------------------------------------------------------------+
    # 🖼️ 4. 图像与窗口尺寸（Image and Window Dimensions）
    IMG_WIDTH = 1920
    IMG_HEIGHT = 1080
    WINDOW_WIDTH = IMG_WIDTH // 3
    WINDOW_HEIGHT = IMG_HEIGHT // 3
    WINDOW_TITLE = "骰子识别与姿态估计"
    ORIGIN_IMAGE_SOURCE = "image"
    # +-----------------------------------------------------------------+


    # +-----------------------------------------------------------------+
    # 🎛️ 5. 图像处理参数（Image Processing Parameters）

    # process_frame 函数
    BOX_OFFSET_VALUE = 5
    CONTOURAREA_THRESHOLD = 20000
    USE_HSV_FLAG = False
    # +-----------------------------------------------------------------+



    # +-----------------------------------------------------------------+
    # 🧮 6. 投影与绘图参数（Drawing Parameters）

    # draw_projections 函数
    OUT_ERROR_FLAG = False

    # display_pose_info 函数
    OUT_BEST_POSE_FLAG = False

    # project_and_draw 函数 （投影线条粗细）
    # REPOJECT_LINE_THICKNESS = 10

    # draw_error_text 函数 （重投影误差文本 绘制参数）
    TEXT_FONT = cv2.FONT_HERSHEY_SIMPLEX
    TEXT_FONT_SCALE = 3.5
    TEXT_FONT_THICKNESS = 10
    TEXT_LINE_HEIGHT = 100
    TOP_X, TOP_Y = 25, 100

    # draw_pose_text 函数 (最优姿态结果 显示参数）
    TEXT_FONT_SCALE_FACTOR = 0.35
    TEXT_LINE_HEIGHT_FACTOR = 0.5
    BOTTOM_X, BOTTOM_Y = 25, 100

    # draw_yolo_datas 函数（YOLO 结果显示参数）
    # YOLO_LINE_THICKNESS = 10
    YOLO_FONT = cv2.FONT_HERSHEY_SIMPLEX
    YOLO_FONT_SCALE = 2
    YOLO_FONT_THICKNESS = 10

    NORMAL_LINE_THICKNESS = 10

    PNP_FRAME_COLOR = (0, 255, 0)
    BA_FRAME_COLOR = (255, 0, 0)
    POSENET_FRAME_COLOR = (255, 204, 51)
    FUSION_FRAME_COLOR = (0, 0, 255)
    # +-----------------------------------------------------------------+


    # +-----------------------------------------------------------------+
    # 🧮 7. Bundle Adjustment 参数（优化算法）
    # bundle_adjustment 函数
    HUBER_LOSS_VALUE = 5.0
    # +-----------------------------------------------------------------+
